import os
from tkinter import Tk, Label, Button, Entry, filedialog, messagebox
from PyPDF2 import PdfReader, PdfWriter

def select_folder():
    """打开文件夹选择对话框"""
    folder_path = filedialog.askdirectory()
    if folder_path:
        folder_entry.delete(0, "end")
        folder_entry.insert(0, folder_path)

def merge_pdfs():
    """提取 PDF 前 n 页并合并"""
    folder_path = folder_entry.get()
    n_pages = pages_entry.get()

    if not folder_path:
        messagebox.showerror("错误", "请选择一个文件夹！")
        return

    if not n_pages.isdigit():
        messagebox.showerror("错误", "请输入有效的页数！")
        return

    n_pages = int(n_pages)
    output_pdf = PdfWriter()

    try:
        # 遍历文件夹内的所有 PDF 文件
        for file_name in os.listdir(folder_path):
            if file_name.endswith(".pdf"):
                file_path = os.path.join(folder_path, file_name)
                reader = PdfReader(file_path)

                # 提取前 n 页
                for page_num in range(min(n_pages, len(reader.pages))):
                    output_pdf.add_page(reader.pages[page_num])

        # 保存合并后的 PDF
        parent_folder = os.path.dirname(folder_path)
        output_path = os.path.join(parent_folder, "merged_output.pdf")
        with open(output_path, "wb") as output_file:
            output_pdf.write(output_file)

        messagebox.showinfo("成功", f"PDF 合并完成！文件已保存到：\n{output_path}")
    except Exception as e:
        messagebox.showerror("错误", f"发生错误：{e}")

# 创建主窗口
root = Tk()
root.title("PDF 合并工具  献给我姐杨莎莎")

# 文件夹选择
Label(root, text="选择文件夹：").grid(row=0, column=0, padx=10, pady=10)
folder_entry = Entry(root, width=40)
folder_entry.grid(row=0, column=1, padx=10, pady=10)
Button(root, text="浏览", command=select_folder).grid(row=0, column=2, padx=10, pady=10)

# 输入页数
Label(root, text="提取页数 (n)：").grid(row=1, column=0, padx=10, pady=10)
pages_entry = Entry(root, width=10)
pages_entry.grid(row=1, column=1, padx=10, pady=10)

# 合并按钮
Button(root, text="合并 PDF", command=merge_pdfs).grid(row=2, column=0, columnspan=3, pady=20)

# 运行主循环
root.mainloop()